import cv2
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial.distance import cdist


def auto_dense_stripe_roi(img,
                          win_size=32,
                          stride=16,
                          thres_percentile=95,
                          margin=5):
    """
    自动检测图像中最密集条纹区域，返回 ROI 框 (x, y, w, h)。
    img : 8-bit 灰度图
    若未检测到，打印提示并返回 None
    """
    h, w = img.shape[:2]

    # 1) 计算梯度能量图
    gx = cv2.Sobel(img, cv2.CV_64F, 1, 0, ksize=3)
    gy = cv2.Sobel(img, cv2.CV_64F, 0, 1, ksize=3)
    grad_mag = np.sqrt(gx ** 2 + gy ** 2)

    # 2) 滑动窗口统计每个块的能量均值与方差
    idx_y, idx_x = np.mgrid[0:h - win_size:stride, 0:w - win_size:stride]
    idx_y = idx_y.ravel()
    idx_x = idx_x.ravel()
    score = np.empty(len(idx_x))

    for k, (y0, x0) in enumerate(zip(idx_y, idx_x)):
        patch = grad_mag[y0:y0 + win_size, x0:x0 + win_size]
        score[k] = np.mean(patch) * np.std(patch)  # 均值×方差，突出条纹

    # 3) 自适应阈值
    th = np.percentile(score, thres_percentile)
    mask = score >= th
    if not np.any(mask):
        print("未检测到明显的条纹密集区域")
        return None

    # 4) 由满足阈值的窗口中心点求最小外接矩形
    centers = np.column_stack((idx_x[mask] + win_size // 2,
                               idx_y[mask] + win_size // 2))
    x_min, y_min = centers.min(axis=0)
    x_max, y_max = centers.max(axis=0)

    # 5) 加边距并限制在图像边界内
    x = max(0, int(x_min - margin))
    y = max(0, int(y_min - margin))
    roi_w = min(w - x, int(x_max - x_min + 2 * margin))
    roi_h = min(h - y, int(y_max - y_min + 2 * margin))

    roi = (x, y, roi_w, roi_h)

    # 6) 可视化
    vis = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
    cv2.rectangle(vis, (x, y), (x + roi_w, y + roi_h), (0, 0, 255), 2)
    plt.figure(figsize=(5, 5))
    plt.title('Auto-detected dense stripe ROI')
    plt.imshow(vis[..., ::-1])
    plt.axis('off')
    plt.show()

    return roi


def brightest_symmetric_spot(image_path, roi=None,
                             inner_radius=15, outer_frac=0.4,
                             thres_rel=0.2, pixel_tol=2):
    """
    仅保留最亮的一对中心对称亮点，返回重建条纹
    """
    # 1. 读图 & ROI
    img = cv2.imread(image_path, 0)
    if img is None:
        raise FileNotFoundError(image_path)
    if roi is None:
        roi = auto_dense_stripe_roi(img)
        # roi = cv2.selectROI('select ROI', img)
        # cv2.destroyAllWindows()

    x, y, w, h = roi
    roi_img = img[y:y + h, x:x + w]

    # 2. FFT + 中心化
    f = np.fft.fft2(roi_img)
    fshift = np.fft.fftshift(f)
    mag = np.abs(fshift)

    # 3. 环形掩码
    cy, cx = h // 2, w // 2
    Y, X = np.ogrid[:h, :w]
    dist = np.sqrt((X - cx) ** 2 + (Y - cy) ** 2)
    outer_radius = min(cx, cy) * outer_frac
    mask = (dist > inner_radius) & (dist < outer_radius)
    mag_band = mag * mask

    # 4. 二值化
    mag_norm = (mag_band - mag_band.min()) / (mag_band.max() - mag_band.min())
    _, bw = cv2.threshold(mag_norm.astype(np.float32),
                          thres_rel, 1, cv2.THRESH_BINARY)

    # 5. 连通域 → 亮点坐标
    bw_u8 = (bw * 255).astype(np.uint8)
    n, labels, stats, centroids = cv2.connectedComponentsWithStats(bw_u8)
    areas = stats[:, cv2.CC_STAT_AREA]
    idx = np.argsort(areas)[::-1][1:]  # 去掉背景
    pts = centroids[idx]

    # 5.5 去掉离中心太近的亮点（含原点）
    dists_to_center = cdist(pts, [[cx, cy]])
    pts = pts[dists_to_center[:, 0] > inner_radius]  # 关键过滤
    if len(pts) < 2:
        print('过滤中心后未找到足够亮点！')
        return None, None

    # 6. 对称配对
    center = np.array([cx, cy])
    pts_centered = pts - center
    pairs = []
    used = set()
    for i, p in enumerate(pts_centered):
        if i in used:
            continue
        p_sym = -p
        dists = cdist([p_sym], pts_centered)[0]
        j = np.argmin(dists)
        if dists[j] < pixel_tol and j not in used:
            pairs.append((pts[i], pts[j]))
            used.update([i, j])

    if not pairs:
        print('未检测到对称亮点！')
        return None, None

    # 7. 仅保留“最亮”的一对
    # 以两个点幅值之和最大作为“最亮”标准
    def pair_intensity(pair):
        p1, p2 = pair
        return mag[int(p1[1]), int(p1[0])] + mag[int(p2[1]), int(p2[0])]

    brightest_pair = max(pairs, key=pair_intensity)
    p1, p2 = brightest_pair

    # 8. 生成频域掩膜
    mask_spots = np.zeros_like(mag, dtype=np.uint8)
    for (xi, yi) in (p1, p2):
        cv2.circle(mask_spots, (int(round(xi)), int(round(yi))), 3, 1, -1)

    f_masked = fshift * mask_spots

    # 9. 反变换
    img_back = np.real(np.fft.ifft2(np.fft.ifftshift(f_masked)))

    # ===== 可视化 =====
    plt.figure(figsize=(12, 4))
    plt.subplot(131)
    plt.title('ROI (original)')
    plt.imshow(roi_img, cmap='gray')
    plt.axis('off')

    plt.subplot(132)
    plt.title('FFT mask (brightest pair)')
    plt.imshow(mask_spots * 255, cmap='gray')
    plt.axis('off')

    plt.subplot(133)
    plt.title('Reconstructed stripes (IFFT)')
    plt.imshow(img_back, cmap='gray')
    # 标出亮点
    for (xi, yi) in (p1, p2):
        plt.plot(xi, yi, 'ro', markersize=6)
    plt.axis('off')
    plt.tight_layout()
    plt.show()

    return brightest_pair, img_back


# ---------------- demo ----------------
if __name__ == '__main__':
    pair, rec = brightest_symmetric_spot('../img/new/0006.tif')
    if pair is not None:
        (x1, y1), (x2, y2) = pair
        print(f'最亮对称亮点：({x1:.1f}, {y1:.1f}) ↔ ({x2:.1f}, {y2:.1f})')
